# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/18 11:17 
@Author  : ZhangShenao 
@File    : 4.使用VectorstoreIndexCreator向量存储检索器.py 
@Desc    : 使用VectorstoreIndexCreator,直接使用向量数据库作为Retriever来进行检索

VectorstoreIndexCreator内部整合了Vectorstore、Embedding、TextSplitter和DocumentLoader等组件
封装了向量检索的整个过程
也可以手动替换这些工具
"""
import os

import dotenv
from langchain.embeddings import OpenAIEmbeddings
from langchain.indexes import VectorstoreIndexCreator
from langchain_community.document_loaders import TextLoader
from langchain_openai import ChatOpenAI

# 加载环境变量
dotenv.load_dotenv()

# 创建LLM
llm = ChatOpenAI(model_name='gpt-3.5-turbo', openai_api_base=os.getenv('OPENAI_API_BASE'))

# 创建Embedding模型
embeddings = OpenAIEmbeddings(
    model='text-embedding-3-small',
    dimensions=1024,  # assign dimensions to 1024
    openai_api_base=os.getenv('OPENAI_API_BASE')
)

# 使用TextLoader,加载文档
loader = TextLoader(file_path='./docs/花语大全.txt')

# 通过VectorstoreIndexCreator来创建索引
# 并在索引的query方法中,通过vectorstore类的as_retriever方法
# 把向量数据库(Vector Store)直接作为检索器,来完成检索任务
# VectorstoreIndexCreator内部整合了Vectorstore、Embedding、TextSplitter和DocumentLoader等组件
# 也可以手动替换这些工具
index = VectorstoreIndexCreator(embedding=embeddings).from_loaders([loader])

# 执行搜索,打印搜索结果
# VectorstoreIndexCreator内部会执行RetrievalQAChain,直接使用向量数据库作为Retriever
print(index.query(question='玫瑰花的花语是什么？', llm=llm))
